syntax = "proto3";

package tensorflow.data;

import "tensorflow/core/data/service/common.proto";
import "tensorflow/core/protobuf/data_service.proto";

// Message representing journaled dispatcher metadata updates. When we apply
// one of these changes to the dispatcher's in-memory state, we also write an
// Update message to the journal.
// Next tag: 17
message Update {
  oneof update_type {
    RegisterDatasetUpdate register_dataset = 1;
    RegisterWorkerUpdate register_worker = 5;
    CreateJobUpdate create_job = 14;
    CreateIterationUpdate create_iteration = 2;
    ProduceSplitUpdate produce_split = 8;
    AcquireIterationClientUpdate acquire_iteration_client = 6;
    ReleaseIterationClientUpdate release_iteration_client = 7;
    GarbageCollectIterationUpdate garbage_collect_iteration = 12;
    RemoveTaskUpdate remove_task = 11;
    CreatePendingTaskUpdate create_pending_task = 9;
    ClientHeartbeatUpdate client_heartbeat = 10;
    CreateTaskUpdate create_task = 3;
    FinishTaskUpdate finish_task = 4;
    SnapshotUpdate snapshot = 15;
    CompressionDisabledAtRuntimeUpdate compression_disabled_at_runtime = 16;
  }
  reserved 13;
}

// Next tag: 5
message RegisterDatasetUpdate {
  string dataset_id = 1;
  uint64 fingerprint = 2;
  DataServiceMetadata metadata = 3;
  bool dedupe_by_dataset_id = 4;
}

// Next tag: 6
message RegisterWorkerUpdate {
  string worker_address = 1;
  repeated DataTransferServerInfo transfer_servers = 5;
  repeated string worker_tags = 3;
  int64 worker_uid = 4;
  reserved 2;
}

// Next tag: 9
message CreateJobUpdate {
  int64 job_id = 1;
  string job_name = 2;
  string dataset_id = 3;
  ProcessingModeDef processing_mode_def = 4;
  // Optional number of consumers. If set, the iteration's tasks will provide
  // their elements to consumers round-robin.
  oneof optional_num_consumers {
    int64 num_consumers = 6;
  }
  // Specifies which workers the client of this iteration reads from.
  TargetWorkers target_workers = 7;
  // True if cross-trainer cache is enabled.
  bool use_cross_trainer_cache = 8;
}

// Next tag: 5
message CreateIterationUpdate {
  int64 iteration_id = 1;
  int64 job_id = 2;
  int64 repetition = 3;
  int64 num_split_providers = 4;
}

// Next tag: 5
message ProduceSplitUpdate {
  int64 iteration_id = 1;
  int64 repetition = 2;
  int64 split_provider_index = 4;
  // Whether the split provider reached its end.
  bool finished = 3;
}

// Next tag: 3
message AcquireIterationClientUpdate {
  int64 iteration_id = 1;
  int64 iteration_client_id = 2;
}

// Next tag: 3
message ReleaseIterationClientUpdate {
  int64 iteration_client_id = 1;
  // The time when the client was released, measured in microseconds since the
  // epoch.
  int64 time_micros = 2;
}

// Next tag: 2
message GarbageCollectIterationUpdate {
  int64 iteration_id = 1;
}

// Next tag: 2
message RemoveTaskUpdate {
  int64 task_id = 1;
}

// Indicates that a client failed to block before reaching the target round.
// Next tag: 2
message TaskRejected {
  // A new target round to try adding the task in.
  int64 new_target_round = 1;
}

// Updates dispatcher state based on a client heartbeat.
// Next tag: 4
message ClientHeartbeatUpdate {
  int64 iteration_client_id = 1;
  bool task_accepted = 2;
  TaskRejected task_rejected = 3;
}

// Next tag: 9
message CreatePendingTaskUpdate {
  int64 task_id = 1;
  int64 iteration_id = 2;
  string worker_address = 3;
  repeated DataTransferServerInfo transfer_servers = 8;
  repeated string worker_tags = 6;
  int64 worker_uid = 7;
  int64 starting_round = 5;
  reserved 4;
}

// Next tag: 10
message CreateTaskUpdate {
  reserved 3, 5;
  int64 task_id = 1;
  int64 iteration_id = 2;
  string worker_address = 4;
  repeated DataTransferServerInfo transfer_servers = 9;
  repeated string worker_tags = 7;
  int64 worker_uid = 8;
  reserved 6;
}

// Next tag: 2
message FinishTaskUpdate {
  int64 task_id = 1;
}

// Next tag: 2
message SnapshotUpdate {
  string path = 1;
}

// Next tag: 3
message CompressionDisabledAtRuntimeUpdate {
  string dataset_id = 1;
  bool compression_disabled = 2;
}
